import json
import requests

LLAMA_API_KEY = 'oWSxh2PDYE2veMM6l0w7aphH'
LLAMA_SECRET_KEY = 'J4WIh7eSvpVyDJxyvGoJ2Xqr1bTmgE7K'

class LlamaChat:
    def __init__(self, api_key=LLAMA_API_KEY, secret_key=LLAMA_SECRET_KEY):
        self.api_key = api_key
        self.secret_key = secret_key
        self.access_token = None

    # 获取Token
    def get_token(self):
        url = f"https://aip.baidubce.com/oauth/2.0/token?grant_type=client_credentials&client_id={self.api_key}&client_secret={self.secret_key}"
        headers = {
            'Content-Type': 'application/json',
            'Accept': 'application/json'
        }
        payload = json.dumps("")

        response = requests.request("POST", url, headers=headers, data=payload)
        try:
            self.access_token = response.json().get("access_token")
        except json.JSONDecodeError as e:
            print(f"Error decoding JSON from token response: {e}")
            self.access_token = None
        return self.access_token

    # 调用模型
    def do_chat(self, message):
        if not self.access_token:
            self.get_token() 

        url = f"https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/qianfan_chinese_llama_2_70b?access_token={self.access_token}"
        headers = {
            'Content-Type': 'application/json'
        }
        payload = json.dumps({
            "messages": [
                {
                    "role": "user",
                    "content": message + "(注：如果用户输入有给出参考网址，那么回答时可搜索参考该网址的内容并在回答中体现)"
                }
            ],
            "response_format": "json_object",
            "stream": False,
            "enable_trace": True
        })

        response = requests.request("POST", url, headers=headers, data=payload)
        try:
            context = response.json()
        except json.JSONDecodeError as e:
            print(f"Error decoding JSON from chat response: {e}")
            return {"error": "格式解析错误，请重试"}

        try:
            result = context['result']
        except (KeyError, json.JSONDecodeError) as e:
            print(f"Error processing chat response: {e}")
            return {"error": "格式解析错误，请重试"}

        return result

if __name__ == '__main__':
    chat_bot = LlamaChat()
    response = chat_bot.do_chat("你好")
    print(response)  
